import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from python_ai.common.xcommon import sep

pd.set_option('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', 1000, 'display.expand_frame_repr', False)
plt.rcParams['font.sans-serif'] = ['Simhei']
plt.rcParams['axes.unicode_minus'] = False

plt.figure(figsize=[12, 6])
spr = 1
spc = 2
spn = 0

df = pd.read_csv('drink.txt')

from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
k = [2, 3, 4, 5, 6]
sse = []  # 误差平方和
si = []  # 轮廓系数
for i in k:
    model = KMeans(n_clusters=i)
    model.fit(df)
    sse.append(model.inertia_)
    h = model.predict(df)
    si.append(silhouette_score(df, h))

spn += 1
plt.subplot(spr, spc, spn)
plt.plot(k, sse)
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(k, si)

# Finally show all drawings.
plt.show()